#include <torch/torch.h>

#include <iostream>
#include <algorithm>

#include "hashencoder.h"

using namespace std;

void main(char* args)
{
	cout << "hello world" << endl;

	uint32_t log2_hashmap_size = 19;
	uint32_t num_levels = 16;
	uint32_t level_dim = 2;
	double base_resolution = 16;
	double finest_resolution = 512;
	double max_params = pow(2, log2_hashmap_size);
	uint32_t offsets[17] = { 0 }; uint32_t offset = 0;
	at::Tensor levelEmbedding[16] = {};
	for (int i = 0; i < num_levels; i++) 
	{
		offsets[i] = offset;
		offset += (uint32_t)max_params;
		at::Tensor local = torch::zeros((int)max_params).unsqueeze_(1);
		torch::nn::init::constant_(local, 0.1 * i);
		local = local.repeat({ 1, 2 });
		levelEmbedding[i] = local;
	}
	offsets[16] = offset;
	at::Tensor embeddings = torch::cat(levelEmbedding, 0).cuda();
	cout << embeddings.sizes() << endl;
	at::Tensor gOffsets = torch::from_blob(offsets, 17, torch::kInt32).cuda();
	double b = exp((log(finest_resolution) - log(base_resolution)) / (num_levels - 1));
	float min[3] = { -2, -2, -2 }; float max[3] = { 2, 2, 2 };
	at::Tensor box_min = torch::from_blob(min, 3, torch::kFloat32).cuda();
	at::Tensor box_max = torch::from_blob(max, 3, torch::kFloat32).cuda();
	at::Tensor inputs = torch::empty({ 2, 3 }, torch::kFloat32).cuda();
	at::Tensor outputs = torch::empty({ 2, 32 }, torch::kFloat32).cuda();
	float xyz1[3] = { 0.48131, 1.0711, 1.5933 };
	float xyz2[3] = { 0.49667, 1.0839, 1.5584 };
	inputs[0] = torch::from_blob(xyz1, 3, torch::kFloat32).cuda();
	inputs[1] = torch::from_blob(xyz2, 3, torch::kFloat32).cuda();
	hash_encoder_forward(&inputs, &embeddings, &gOffsets, &box_min, &box_max, base_resolution, b, &outputs, nullptr);
	at::Tensor view = outputs.detach().cpu().view(-1);
	cout << view << endl;
	at::Tensor grad_embeddings = torch::zeros_like(embeddings).cuda();
	at::Tensor grad = torch::empty({ 2, 32 }, torch::kFloat32).cuda();
	torch::nn::init::constant_(grad, 15);
	hash_encoder_backward(&grad, &inputs, &embeddings, &gOffsets, &box_min, &box_max, base_resolution, b, &grad_embeddings);
	embeddings += grad_embeddings;
	outputs = outputs.cuda();
	hash_encoder_forward(&inputs, &embeddings, &gOffsets, &box_min, &box_max, base_resolution, b, &outputs, nullptr);
	view = outputs.detach().cpu().view(-1);
	cout << view << endl;

}